import os

from options import parse_args


def face_templates_count(lfw_all_path, func, template):
    for name in os.listdir(lfw_all_path):
        name_path = os.path.join(lfw_all_path, name)
        num = len(os.listdir(name_path))
        if num >= 5:
            func(template.format(name, num))


def main():
    args = parse_args()
    lfw_all_path = os.path.join(args.root_dir, 'lfw_all')
    args.save = True
    if args.save:
        with open('multiple_faces_statistics.dat', 'w') as f:
            face_templates_count(lfw_all_path, f.write, '{},{}\n')
    else:
        face_templates_count(lfw_all_path, print, '{},{}')


if __name__ == '__main__':
    main()
